import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# Q 网络

class QNetwork(nn.Module):
    def __init__(self, width, height, action_size):
        super(QNetwork, self).__init__()
        self.conv1 = nn.Conv2d(2, 16, kernel_size=3, stride=1, padding=1)  # 输入通道改为2
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.flatten_size = 32 * width * height  # 移除池化层后的尺寸
        self.fc1 = nn.Linear(self.flatten_size, 128)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, action_size)

    def forward(self, x):
        x = self.bn1(torch.relu(self.conv1(x)))
        x = self.bn2(torch.relu(self.conv2(x)))
        x = x.view(-1, self.flatten_size)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x


# 经验回放缓冲区
class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)

    def __len__(self):
        return len(self.buffer)

    def clear(self):
        self.buffer.clear()


class SumTree:
    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = np.zeros(2 * capacity - 1)  # SumTree的节点数
        self.data = np.zeros(capacity, dtype=object)  # 存储经验
        self.write_pointer = 0
        self.n_entries = 0

    def _propagate(self, idx, change):
        parent = (idx - 1) // 2
        self.tree[parent] += change
        if parent != 0:
            self._propagate(parent, change)

    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1
        if left >= len(self.tree):
            return idx
        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def total(self):
        return self.tree[0]

    def add(self, priority, data):
        idx = self.write_pointer + self.capacity - 1
        self.data[self.write_pointer] = data
        self.update(idx, priority)
        self.write_pointer = (self.write_pointer + 1) % self.capacity
        if self.n_entries < self.capacity:
            self.n_entries += 1

    def update(self, idx, priority):
        change = priority - self.tree[idx]
        self.tree[idx] = priority
        self._propagate(idx, change)

    def get(self, s):
        idx = self._retrieve(0, s)
        data_idx = idx - self.capacity + 1
        return idx, self.tree[idx], self.data[data_idx]


class PrioritizedReplayBuffer:
    def __init__(self, capacity, alpha=0.6, beta=0.4, beta_increment=0.001):
        self.tree = SumTree(capacity)
        self.capacity = capacity
        self.alpha = alpha  # 优先级影响因子
        self.beta = beta  # 初始重要性采样权重因子
        self.beta_increment = beta_increment  # beta增量
        self.max_priority = 1.0  # 初始化最大优先级
        self.epsilon = 1e-6  # 避免优先级为0

    def push(self, state, action, reward, next_state, done):
        # 创建经验
        experience = (state, action, reward, next_state, done)
        # 使用最大优先级存储新经验
        priority = self.max_priority
        self.tree.add(priority, experience)

    def sample(self, batch_size):
        batch = []
        idxs = []
        priorities = []
        segment = self.tree.total() / batch_size
        weights = []

        for i in range(batch_size):
            s = np.random.uniform(i * segment, (i + 1) * segment)
            idx, priority, data = self.tree.get(s)
            batch.append(data)
            idxs.append(idx)
            priorities.append(priority)
            # 计算重要性采样权重
            prob = priority / self.tree.total()
            weights.append((self.tree.n_entries * prob) ** (-self.beta))

        # 归一化权重
        weights = np.array(weights) / np.max(weights)
        # 更新beta
        self.beta = min(1.0, self.beta + self.beta_increment)
        return batch, idxs, weights

    def update_priorities(self, idxs, td_errors):
        for idx, td_error in zip(idxs, td_errors):
            priority = (abs(td_error) + self.epsilon) ** self.alpha
            self.tree.update(idx, priority)
            self.max_priority = max(self.max_priority, priority)

    def __len__(self):
        return self.tree.n_entries

    def clear(self):
        self.tree = SumTree(self.capacity)
        self.write_pointer = 0
        self.n_entries = 0


# DQN 代理
class DQNAgent:
    def __init__(self, shape, action_size, gamma=0.99, epsilon=1, epsilon_min=0.05, epsilon_decay=0.99,
                 learning_rate=0.002, batch_size=32, memory_size=50000):
        self.width = shape[1]
        self.height = shape[0]
        self.action_size = action_size  # 动态动作数量
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min
        self.epsilon_decay = epsilon_decay
        self.batch_size = batch_size
        # self.memory = ReplayBuffer(memory_size)
        self.memory = PrioritizedReplayBuffer(memory_size, alpha=0.6, beta=0.4)
        # 初始化 Q 网络和目标网络
        self.q_network = QNetwork(shape[1], shape[0], action_size).to(device)
        self.target_network = QNetwork(shape[1], shape[0], action_size).to(device)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=learning_rate)

    def reset_epsilon(self):
        self.epsilon = 1.0  # 每次重置地图时恢复探索

    # 限制机器人不想越界方向运动
    def is_over(self, action, env, q_net_args):
        if not env.take_actions(list(env.elements.keys())[0], int(action), sim_=True):
            q_net_args[:, action] = -10000  # 让机器人不选择越出地图边界的选择
            return self.is_over(q_net_args.argmax().item(), env, q_net_args)
        else:
            return action

    def greedy_policy(self, bot, env):
        step_choice = [[0, -1], [0, 1], [-1, 0], [1, 0], [1, -1], [-1, -1], [1, 1], [-1, 1]]
        score = []
        quality_scores = np.array(env.grid_map * env.Gaussian_mass_map)
        quality_scores_map = quality_scores / (np.max(quality_scores) + 1e-8)
        max_concentration_pos = np.unravel_index(np.argmax(quality_scores_map), quality_scores_map.shape)
        for choice in step_choice:
            next_pos = np.array(bot.current_grid_pos) + choice
            if 0 <= next_pos[0] < env.grid_map.shape[0] and 0 <= next_pos[1] < env.grid_map.shape[1]:
                distance = np.linalg.norm(np.array(max_concentration_pos) - next_pos)
                score.append(distance)
            else:
                score.append(float('inf'))
        return np.argmin(score)

    # def greedy_policy(self, bot, target_source_visual):
    #     step_choice = [[0, -1], [0, 1], [-1, 0], [1, 0], [1, -1], [-1, -1], [1, 1], [-1, 1]]
    #     score = []
    #     # action = list(self.action_dict.keys())[action]
    #     for index, choice in enumerate(step_choice):
    #         next_pos = (np.array(bot.current_grid_pos) + choice)
    #         Prior_distance = np.linalg.norm(np.array(target_source_visual) - next_pos)
    #         score.append(Prior_distance)
    #     probabilities = np.array(score) / np.sum(score)
    #     return np.argmin(probabilities)

    def take_action(self, state, env, element_name, target_source_visual=None):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            q_values = self.q_network(state)
            # 随机选择动作
            print(self.epsilon)
            if random.random() < self.epsilon:
                # 基于先验地图的加权随机探索,更靠近于先验地图的概率会更高
                action = self.greedy_policy(env.element_collection[element_name], env)
            else:
                # print("Q值动作")
                # print(q_values)
                action = q_values.argmax().item()
        return action

    def store_experience(self, state, action, reward, next_state, done):
        self.memory.push(state, action, reward, next_state, done)

    def train(self):
        if len(self.memory) < self.batch_size:
            return
        # batch = self.memory.sample(self.batch_size)
        batch, idxs, weights = self.memory.sample(self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(np.array(states)).to(device)
        actions = torch.LongTensor(actions).to(device)
        rewards = torch.FloatTensor(rewards).to(device)
        next_states = torch.FloatTensor(np.array(next_states)).to(device)
        dones = torch.FloatTensor(dones).to(device)
        weights = torch.FloatTensor(weights).to(device)  # Convert weights to tensor and move to device

        q_values = self.q_network(states).gather(1, actions.unsqueeze(1)).squeeze(1)
        with torch.no_grad():
            next_actions = self.q_network(next_states).argmax(1)  # 当前网络选择动作
            next_q_values = self.target_network(next_states).gather(1, next_actions.unsqueeze(1)).squeeze(1)
            target_q_values = rewards + (1 - dones) * self.gamma * next_q_values
            # next_q_values = self.target_network(next_states).max(1)[0]
            # target_q_values = rewards + (1 - dones) * self.gamma * next_q_values

        # 计算TD误差
        td_errors = (q_values - target_q_values).abs().detach().cpu().numpy()
        # print("TD errors:", td_errors)

        # 更新优先级
        self.memory.update_priorities(idxs, td_errors)

        # 计算带权重的损失
        loss = (weights * nn.MSELoss(reduction='none')(q_values, target_q_values)).mean()
        # loss = nn.MSELoss()(q_values, target_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

    def update_target_network(self):
        self.target_network.load_state_dict(self.q_network.state_dict())

    def model_save(self):
        try:
            torch.save(self.q_network.state_dict(), 'net_eval.pkl')
            torch.save(self.target_network.state_dict(), 'net_target.pkl')
            print("successfully save")
        except Exception:
            print('save fail')

    def best_model_save(self):
        try:
            torch.save(self.q_network.state_dict(), f'best_net_eval.pkl')
            torch.save(self.target_network.state_dict(), 'best_net_target.pkl')
            print("successfully save")
        except Exception:
            print('save fail')

    def model_read(self):
        try:
            self.q_network.load_state_dict(torch.load('net_eval.pkl'))
            self.target_network.load_state_dict(torch.load('net_target.pkl'))
            print("successfully read")
        except Exception:
            print('read fail')
